"""
简化的黄金和美元指数数据获取脚本
"""

import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import os
import yfinance as yf


def generate_mock_gold_data(start_date="2014-01-01", end_date=None):
    """生成模拟黄金数据"""
    if end_date is None:
        end_date = datetime.now().strftime("%Y-%m-%d")
    
    print("生成模拟黄金数据...")
    
    # 创建日期范围（只包含工作日）
    dates = pd.date_range(start=start_date, end=end_date, freq='B')  # B = Business day
    
    # 生成模拟价格数据
    np.random.seed(42)  # 确保可重复性
    n_days = len(dates)
    
    # 初始价格（美元/盎司）
    initial_price = 1200.0
    
    # 生成随机游走价格
    returns = np.random.normal(0.0005, 0.015, n_days)  # 日收益率
    prices = [initial_price]
    
    for i in range(1, n_days):
        new_price = prices[-1] * (1 + returns[i])
        prices.append(max(new_price, 100))  # 确保价格不会太低
    
    prices = np.array(prices)
    
    # 生成OHLC数据
    data = pd.DataFrame({
        'date': dates,
        'close': prices
    })
    
    # 生成开盘价
    data['open'] = np.roll(data['close'], 1)
    data.loc[0, 'open'] = initial_price
    data['open'] = data['open'] * (1 + np.random.normal(0, 0.005, n_days))
    
    # 生成最高价和最低价
    data['high'] = np.maximum(data['open'], data['close']) * (1 + np.abs(np.random.normal(0, 0.01, n_days)))
    data['low'] = np.minimum(data['open'], data['close']) * (1 - np.abs(np.random.normal(0, 0.01, n_days)))
    
    # 生成成交量
    data['volume'] = np.random.randint(100000, 500000, n_days)
    
    # 添加标识
    data['symbol'] = 'GOLD'
    data['data_source'] = 'MOCK'
    
    return data


def generate_mock_usd_data(start_date="2014-01-01", end_date=None):
    """生成模拟美元指数数据"""
    if end_date is None:
        end_date = datetime.now().strftime("%Y-%m-%d")
    
    print("生成模拟美元指数数据...")
    
    # 创建日期范围（只包含工作日）
    dates = pd.date_range(start=start_date, end=end_date, freq='B')
    
    # 生成模拟价格数据
    np.random.seed(123)  # 确保可重复性
    n_days = len(dates)
    
    # 初始美元指数
    initial_index = 80.0
    
    # 生成随机游走价格
    returns = np.random.normal(0.0002, 0.008, n_days)  # 日收益率
    prices = [initial_index]
    
    for i in range(1, n_days):
        new_price = prices[-1] * (1 + returns[i])
        prices.append(max(new_price, 70))  # 确保指数不会太低
    
    prices = np.array(prices)
    
    # 生成OHLC数据
    data = pd.DataFrame({
        'date': dates,
        'close': prices
    })
    
    # 生成开盘价
    data['open'] = np.roll(data['close'], 1)
    data.loc[0, 'open'] = initial_index
    data['open'] = data['open'] * (1 + np.random.normal(0, 0.003, n_days))
    
    # 生成最高价和最低价
    data['high'] = np.maximum(data['open'], data['close']) * (1 + np.abs(np.random.normal(0, 0.005, n_days)))
    data['low'] = np.minimum(data['open'], data['close']) * (1 - np.abs(np.random.normal(0, 0.005, n_days)))
    
    # 生成成交量
    data['volume'] = np.random.randint(50000, 200000, n_days)
    
    # 添加标识
    data['symbol'] = 'USD_INDEX'
    data['data_source'] = 'MOCK'
    
    return data


def save_data(data, filename, directory="data/raw"):
    """保存数据到CSV文件"""
    os.makedirs(directory, exist_ok=True)
    filepath = os.path.join(directory, filename)
    data.to_csv(filepath, index=False, encoding='utf-8-sig')
    print(f"数据已保存到: {filepath}")
    return filepath


def main():
    """主函数"""
    print("=" * 60)
    print("黄金交易预测数据获取程序（模拟数据版本）")
    print("=" * 60)
    
    # 设置时间范围
    start_date = "2014-01-01"
    end_date = datetime.now().strftime("%Y-%m-%d")
    
    print(f"数据时间范围: {start_date} 到 {end_date}")
    print()
    
    # 生成黄金数据
    gold_data = generate_mock_gold_data(start_date, end_date)
    print(f"生成黄金数据: {len(gold_data)} 条记录")
    print(f"时间范围: {gold_data['date'].min()} 到 {gold_data['date'].max()}")
    print(f"价格范围: ${gold_data['close'].min():.2f} - ${gold_data['close'].max():.2f}")
    print()
    
    # 生成美元指数数据
    usd_data = generate_mock_usd_data(start_date, end_date)
    print(f"生成美元指数数据: {len(usd_data)} 条记录")
    print(f"时间范围: {usd_data['date'].min()} 到 {usd_data['date'].max()}")
    print(f"指数范围: {usd_data['close'].min():.2f} - {usd_data['close'].max():.2f}")
    print()
    
    # 合并数据
    combined_data = pd.merge(
        gold_data[['date', 'close']].rename(columns={'close': 'gold_close'}),
        usd_data[['date', 'close']].rename(columns={'close': 'usd_close'}),
        on='date',
        how='inner'
    )
    
    print(f"合并后数据: {len(combined_data)} 条记录")
    print()
    
    # 保存数据
    save_data(gold_data, 'gold_spot_data.csv')
    save_data(usd_data, 'usd_index_data.csv')
    save_data(combined_data, 'combined_data.csv')
    
    # 显示数据样例
    print("数据样例:")
    print(combined_data.head())
    print()
    
    print("数据统计信息:")
    print(combined_data.describe())
    print()
    
    print("数据获取完成！")
    print("=" * 60)


if __name__ == "__main__":
    main()
